import sys
from PyQt5.QtWidgets import QApplication, QMainWindow
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle
import numpy as np

class SelectBox:
    def __init__(self, ax):
        self.ax = ax
        self.press = None
        self.rect = None
        self.selected_points = []

        self.cidmotion = None  # 添加初始化

        self.cidrelease = None  # 添加初始化

    def on_press(self, event):
        if event.button != 1:
            return
        if event.inaxes != self.ax:
            return
        self.press = (event.xdata, event.ydata)
        self.rect = Rectangle((event.xdata, event.ydata), 0, 0, edgecolor='r', facecolor='none', linewidth=2)
        self.ax.add_patch(self.rect)
        self.ax.figure.canvas.draw_idle()

        # 这里添加对鼠标移动和释放事件的连接，并保存它们的连接ID

        self.cidmotion = self.ax.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)

        self.cidrelease = self.ax.figure.canvas.mpl_connect('button_release_event', self.on_release)

    def points_in_bbox(points, bbox):

        """Check which points are within the bounding box."""

        x_min, y_min, x_max, y_max = bbox.extents

        return [(x, y) for x, y in points if x_min <= x <= x_max and y_min <= y <= y_max]

    def on_release(self, event):

        if self.press is None: return

        # Remove the rectangel and disconnect the callbacks

        self.ax.figure.canvas.mpl_disconnect(self.cidmotion)

        self.ax.figure.canvas.mpl_disconnect(self.cidrelease)

        self.ax.patches.remove(self.rect)

        self.press = None

        # Update the rectangle position to exactly cover the selection area

        width = event.xdata - self.press[0]

        height = event.ydata - self.press[1]

        self.rect.set_width(width)

        self.rect.set_height(height)

        self.rect.set_xy((self.press[0], self.press[1]))

        # Assuming you have a list of data points called `data_points`

        bbox = self.rect.get_bbox().transformed(self.ax.transData.inverted())

        selected_points = points_in_bbox(data_points, bbox)

        print("Selected Points:")

        for point in selected_points:
            print(point)

        self.ax.figure.canvas.draw_idle()

    # def on_release(self, event):
    #     if event.button != 1 or self.press is None:
    #         return
    #     if event.inaxes != self.ax:
    #         return
    #     self.ax.figure.canvas.mpl_disconnect(self.cidmotion)
    #     self.ax.figure.canvas.mpl_disconnect(self.cidrelease)
    #
    #     x0, y0 = self.press
    #     x1, y1 = event.xdata, event.ydata
    #
    #     # 确保是矩形
    #     x0, x1 = sorted([x0, x1])
    #     y0, y1 = sorted([y0, y1])
    #
    #     # 假设这里有一个数据点列表，您需要替换为实际数据
    #     # 示例数据点
    #     x_data = np.random.rand(10)
    #     y_data = np.random.rand(10)
    #
    #     # 筛选框选中的点
    #     self.selected_points = [(xi, yi) for xi, yi in zip(x_data, y_data) if x0 <= xi <= x1 and y0 <= yi <= y1]
    #
    #     # 打印选中的点（实际应用中，您可能需要在这里进行高亮显示或其他处理）
    #     print("Selected points:", self.selected_points)
    #
    #     # 移除矩形框
    #     self.rect.remove()
    #     self.rect = None
    #     self.press = None
    #     self.ax.figure.canvas.draw_idle()

    def on_motion(self, event):
        if self.press is None:
            return
        if event.inaxes != self.ax:
            return

        # 更新矩形框的位置和大小
        self.rect.set_width(event.xdata - self.press[0])
        self.rect.set_height(event.ydata - self.press[1])
        self.ax.figure.canvas.draw_idle()

class DualAxisWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.fig = Figure()
        self.ax = self.fig.add_subplot(111)

        # 示例数据点，实际应用中应替换为您的数据
        x_data = np.random.rand(10)
        y_data = np.random.rand(10)
        self.ax.plot(x_data, y_data, 'o')  # 使用散点图以方便高亮显示

        self.select_box = SelectBox(self.ax)

        self.canvas = FigureCanvas(self.fig)
        self.setCentralWidget(self.canvas)

        # 连接事件
        self.cidpress = self.canvas.mpl_connect('button_press_event', self.select_box.on_press)
        self.cidmotion = self.canvas.mpl_connect('motion_notify_event', self.select_box.on_motion)
        self.cidrelease = self.canvas.mpl_connect('button_release_event', self.select_box.on_release)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    mainWin = DualAxisWindow()
    mainWin.show()
    sys.exit(app.exec_())